# 预算优化Agent - 智能预算分析和Agent间协商
from python_a2a import A2AServer, run_server, TaskStatus, TaskState, AgentCard, AgentSkill
import requests
import json
import uuid
import logging
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('BudgetAgent')

class BudgetAgent(A2AServer):
    """预算优化专家 - 支持真正的A2A直接协商"""

    def __init__(self, agent_card):
        super().__init__(agent_card=agent_card)
        
        # 预算分类标准
        self.budget_categories = {
            "accommodation": {"min_ratio": 0.25, "max_ratio": 0.40, "priority": "high"},
            "food": {"min_ratio": 0.20, "max_ratio": 0.35, "priority": "medium"},
            "transportation": {"min_ratio": 0.10, "max_ratio": 0.25, "priority": "high"},
            "activities": {"min_ratio": 0.15, "max_ratio": 0.30, "priority": "medium"},
            "emergency": {"min_ratio": 0.05, "max_ratio": 0.10, "priority": "high"}
        }
        
        # A2A协商支持 - 真正的点对点通信
        self.peer_agents = {
            "weather": "http://localhost:7003",
            "travel": "http://localhost:7004",
            "restaurant": "http://localhost:7005"
        }
        
        # 协商状态管理
        self.active_direct_negotiations = {}  # 直接协商会话
        self.peer_negotiation_history = []    # 对等协商历史

        print(f"💰 BudgetAgent 初始化完成，支持A2A直接协商")

    def handle_task(self, task):
        """处理预算相关任务"""
        message_data = task.message or {}
        content = message_data.get("content", {})
        text = content.get("text", "")
        
        logger.info(f"💰 BudgetAgent收到任务: {text}")
        print(f"💰 BudgetAgent收到任务: {text[:100]}...")
        
        try:
            # 解析请求
            if text.startswith('{'):
                request_data = json.loads(text)
            else:
                request_data = {"query": text, "type": "simple_budget_query"}
            
            # 根据请求类型处理
            if request_data.get("type") == "collaboration_request":
                response = self._handle_collaboration_request(request_data)
            elif request_data.get("type") == "negotiation_request":
                response = self._handle_negotiation_request(request_data)
            else:
                response = self._handle_simple_budget_analysis(request_data)
                
        except Exception as e:
            logger.error(f"❌ BudgetAgent处理错误: {e}")
            response = {
                "agent": "BudgetAgent",
                "error": f"处理预算分析时出错: {str(e)}",
                "message": text
            }

        task.artifacts = [{"parts": [{"type": "text", "text": json.dumps(response, ensure_ascii=False)}]}]
        task.status = TaskStatus(state=TaskState.COMPLETED)
        logger.info("✅ BudgetAgent任务完成")
        
        # 发送数据流信息到UI
        try:
            import requests
            requests.post("http://localhost:8000/api/notify", json={
                "type": "data_flow",
                "step": {
                    "agent": "BudgetAgent",
                    "description": f"预算分析完成 - {request_data.get('type', 'unknown')}",
                    "input": request_data,
                    "output": response,
                    "upstream": "CoordinatorAgent",
                    "downstream": "CoordinatorAgent"
                }
            }, timeout=2)
        except:
            pass
        
        return task

    def _handle_collaboration_request(self, request_data):
        """处理协作请求"""
        user_requirements = request_data.get("user_requirements", {})
        upstream_context = request_data.get("upstream_context", {})
        collaboration_context = request_data.get("collaboration_context", {})
        
        logger.info(f"💰 开始预算协作分析")
        
        # 1. 基础预算分析
        base_budget = self._extract_budget_from_requirements(user_requirements)
        budget_breakdown = self._analyze_budget_breakdown(base_budget, user_requirements)
        
        # 2. 分析上游影响（主要是天气影响）
        weather_impact = upstream_context.get("weather_impact", {})
        if weather_impact:
            budget_breakdown = self._adjust_for_weather_impact(budget_breakdown, weather_impact)
        
        # 3. 检测预算冲突
        conflicts = self._detect_budget_conflicts(budget_breakdown, user_requirements)
        
        # 4. 生成协商方案
        negotiation_points = self._generate_negotiation_points(conflicts, budget_breakdown, user_requirements)
        
        # 5. 如果发现严重冲突，主动发起协商
        if conflicts and any(c["severity"] == "high" for c in conflicts):
            self._initiate_active_negotiation(conflicts, user_requirements)
        
        budget_status = self._get_budget_status(base_budget, budget_breakdown)
        
        # 如果预算不可行，返回拒绝响应
        if budget_status == "impossible":
            response = {
                "type": "collaboration_response",
                "agent": "BudgetAgent",
                "budget_data": {
                    "base_budget": base_budget,
                    "budget_status": "impossible",
                    "error": f"预算{base_budget}欧元严重不足，无法安排任何旅行活动",
                    "min_required_budget": 200,
                    "recommendation": "强烈建议增加预算至少200欧元或考虑取消旅行计划"
                },
                "potential_conflicts": conflicts,
                "confidence_score": 0.95,
                "final_recommendation": "TRAVEL_NOT_RECOMMENDED"  # 明确标记不推荐旅行
            }
        else:
            response = {
                "type": "collaboration_response",
                "agent": "BudgetAgent",
                "budget_data": {
                    "base_budget": base_budget,
                    "optimized_breakdown": budget_breakdown,
                    "total_required": sum(budget_breakdown.values()),
                    "budget_status": budget_status,
                    "flexibility_analysis": self._analyze_budget_flexibility(budget_breakdown, user_requirements)
                },
                "impact_analysis": {
                    "weather_cost_adjustment": weather_impact.get("extra_cost", 0),
                    "cost_optimization_potential": self._calculate_optimization_potential(budget_breakdown),
                    "risk_factors": self._identify_cost_risks(budget_breakdown, user_requirements)
                },
                "negotiation_points": negotiation_points,
                "potential_conflicts": conflicts,
                "confidence_score": 0.88,
                "collaboration_suggestions": self._suggest_collaboration_strategy(conflicts, budget_breakdown)
            }
        
        return response

    def _handle_negotiation_request(self, request_data):
        """处理来自其他Agent的协商请求"""
        sender = request_data.get("sender", "Unknown")
        negotiation_context = request_data.get("negotiation_context", {})
        proposals = request_data.get("proposals", [])
        
        logger.info(f"💰 处理来自{sender}的协商请求")
        
        # 分析每个提案的预算影响
        proposal_analysis = []
        for proposal in proposals:
            analysis = self._analyze_proposal_impact(proposal, negotiation_context)
            proposal_analysis.append(analysis)
        
        # 生成反向提案
        counter_proposals = self._generate_counter_proposals(proposal_analysis, negotiation_context)
        
        # 选择最优方案
        recommended_solution = self._select_optimal_solution(proposal_analysis, counter_proposals)
        
        response = {
            "type": "negotiation_response",
            "agent": "BudgetAgent",
            "sender_ack": sender,
            "proposal_analysis": proposal_analysis,
            "counter_proposals": counter_proposals,
            "recommended_solution": recommended_solution,
            "negotiation_status": "active",
            "next_steps": ["await_sender_response", "finalize_budget_allocation"]
        }
        
        return response

    def _extract_budget_from_requirements(self, user_requirements):
        """从用户需求中提取预算信息"""
        budget_text = user_requirements.get("budget", "") or user_requirements.get("base_budget", "")
        
        # 提取数字
        import re
        numbers = re.findall(r'\d+', str(budget_text))
        if numbers:
            budget_amount = int(numbers[0])
            print(f"💰 检测到用户预算: {budget_amount}")
            return budget_amount
        
        # 如果完全没有预算信息，返回一个明显不足的预算来触发警告
        print("⚠️ 未找到明确的预算信息，使用最低预算触发预算不足警告")
        return 50  # 极低预算，会触发预算不足的处理逻辑

    def _analyze_budget_breakdown(self, base_budget, user_requirements):
        """分析预算分解"""
        preferences = user_requirements.get("preferences", "") or user_requirements.get("interests", [])
        
        # 基础分配
        breakdown = {
            "accommodation": base_budget * 0.30,
            "food": base_budget * 0.25,
            "transportation": base_budget * 0.15,
            "activities": base_budget * 0.25,
            "emergency": base_budget * 0.05
        }
        
        # 根据偏好调整
        if isinstance(preferences, str):
            if "美食" in preferences or "food" in preferences.lower():
                breakdown["food"] += base_budget * 0.05
                breakdown["activities"] -= base_budget * 0.05
            if "历史" in preferences or "museum" in preferences.lower():
                breakdown["activities"] += base_budget * 0.05
                breakdown["transportation"] -= base_budget * 0.05
        
        return breakdown

    def _adjust_for_weather_impact(self, budget_breakdown, weather_impact):
        """根据天气影响调整预算"""
        extra_transport = weather_impact.get("extra_transportation_cost", 0)
        equipment_cost = weather_impact.get("equipment_cost", 0)
        
        if extra_transport > 0:
            budget_breakdown["transportation"] += extra_transport
        
        if equipment_cost > 0:
            budget_breakdown["emergency"] += equipment_cost
            
        return budget_breakdown

    def _detect_budget_conflicts(self, budget_breakdown, user_requirements):
        """检测预算冲突"""
        base_budget = self._extract_budget_from_requirements(user_requirements)
        total_required = sum(budget_breakdown.values())
        
        conflicts = []
        
        # 检查极低预算情况
        min_viable_budget = 200  # 最低可行预算
        if base_budget < min_viable_budget:
            conflicts.append({
                "type": "budget_too_low",
                "severity": "critical",
                "amount": base_budget,
                "min_required": min_viable_budget,
                "shortage": min_viable_budget - base_budget,
                "description": f"预算{base_budget}欧元严重不足，无法安排任何有意义的旅行活动。最低需要{min_viable_budget}欧元。",
                "recommendation": "建议大幅增加预算或考虑取消旅行计划"
            })
            return conflicts  # 极低预算时直接返回，不进行其他检查
        
        if total_required > base_budget:
            shortage = total_required - base_budget
            severity = "critical" if shortage > base_budget * 0.5 else ("high" if shortage > base_budget * 0.2 else "medium")
            
            conflicts.append({
                "type": "budget_shortfall",
                "severity": severity,
                "amount": shortage,
                "percentage": (shortage / base_budget) * 100,
                "description": f"预算不足{shortage:.0f}欧元（{(shortage/base_budget)*100:.1f}%）",
                "affected_categories": self._identify_overbudget_categories(budget_breakdown, base_budget)
            })
        
        # 检查单项预算超标
        for category, amount in budget_breakdown.items():
            max_ratio = self.budget_categories[category]["max_ratio"]
            if amount > base_budget * max_ratio:
                conflicts.append({
                    "type": "category_overbudget",
                    "severity": "medium",
                    "category": category,
                    "amount": amount,
                    "max_recommended": base_budget * max_ratio,
                    "description": f"{category}预算超出推荐范围"
                })
        
        return conflicts

    def _generate_negotiation_points(self, conflicts, budget_breakdown, user_requirements):
        """生成协商要点"""
        negotiation_points = []
        
        for conflict in conflicts:
            if conflict["type"] == "budget_shortfall":
                # 与天气Agent协商
                negotiation_points.append({
                    "target_agent": "WeatherAgent",
                    "topic": "cost_reduction_options",
                    "proposal": "探索降低天气相关成本的方案",
                    "details": {
                        "current_shortage": conflict["amount"],
                        "acceptable_weather_tradeoffs": ["调整日期", "增加室内活动比例"]
                    }
                })
                
                # 与用户协商（通过Coordinator）
                negotiation_points.append({
                    "target_agent": "CoordinatorAgent",
                    "topic": "budget_increase_request",
                    "proposal": f"建议用户增加预算{conflict['amount']:.0f}欧元",
                    "justification": "确保行程质量和安全性",
                    "alternatives": self._generate_budget_alternatives(budget_breakdown, conflict["amount"])
                })
        
        return negotiation_points

    def _generate_budget_alternatives(self, budget_breakdown, shortage):
        """生成预算替代方案"""
        alternatives = []
        
        # 方案1：降低住宿标准
        alt1 = budget_breakdown.copy()
        alt1["accommodation"] *= 0.8
        alternatives.append({
            "name": "经济住宿",
            "description": "选择更经济的住宿选项",
            "savings": budget_breakdown["accommodation"] * 0.2,
            "impact": "住宿舒适度略降，但位置仍便利"
        })
        
        # 方案2：减少活动预算
        alt2 = budget_breakdown.copy()
        alt2["activities"] *= 0.7
        alternatives.append({
            "name": "精选活动",
            "description": "重点体验最核心的活动",
            "savings": budget_breakdown["activities"] * 0.3,
            "impact": "活动数量减少，但质量保持"
        })
        
        return alternatives

    def _initiate_active_negotiation(self, conflicts, user_requirements):
        """主动发起与其他Agent的协商"""
        logger.info("💰 检测到严重预算冲突，主动发起协商")
        
        # 向天气Agent发起协商
        try:
            weather_negotiation = {
                "type": "negotiation_request",
                "sender": "BudgetAgent",
                "negotiation_context": {
                    "issue": "budget_shortfall",
                    "severity": "high",
                    "requested_cost_reduction": conflicts[0]["amount"]
                },
                "proposals": [
                    {
                        "option": "adjust_travel_dates",
                        "description": "调整旅行日期以降低成本",
                        "expected_savings": conflicts[0]["amount"] * 0.6
                    },
                    {
                        "option": "increase_indoor_ratio",
                        "description": "增加室内活动比例，减少交通成本",
                        "expected_savings": conflicts[0]["amount"] * 0.4
                    }
                ]
            }
            
            response = self._send_negotiation_to_agent("weather", weather_negotiation)
            logger.info(f"💰 已向WeatherAgent发起协商，响应: {response.status_code}")
            
        except Exception as e:
            logger.warning(f"💰 主动协商失败: {e}")

    def _send_negotiation_to_agent(self, agent_type, negotiation_data):
        """向指定Agent发送协商请求"""
        if agent_type not in self.agents:
            raise ValueError(f"未知的Agent类型: {agent_type}")
        
        agent_url = self.agents[agent_type]
        payload = {
            "message": {
                "content": {
                    "text": json.dumps(negotiation_data, ensure_ascii=False)
                }
            }
        }
        
        response = requests.post(f"{agent_url}/a2a/tasks/send", json=payload, timeout=30)
        return response

    def _analyze_proposal_impact(self, proposal, context):
        """分析提案的预算影响"""
        option = proposal.get("option", "")
        
        if "adjust_dates" in option:
            return {
                "proposal": proposal,
                "cost_impact": "可能节省10-20%交通和住宿成本",
                "feasibility": "high",
                "user_impact": "medium",
                "recommendation": "acceptable"
            }
        elif "indoor" in option:
            return {
                "proposal": proposal,
                "cost_impact": "节省15-25%交通成本",
                "feasibility": "high", 
                "user_impact": "low",
                "recommendation": "preferred"
            }
        
        return {
            "proposal": proposal,
            "cost_impact": "unknown",
            "feasibility": "medium",
            "user_impact": "medium",
            "recommendation": "needs_analysis"
        }

    def _generate_counter_proposals(self, proposal_analysis, context):
        """生成反向提案"""
        counter_proposals = []
        
        # 基于分析结果生成反提案
        counter_proposals.append({
            "option": "hybrid_cost_optimization",
            "description": "组合方案：轻微调整日期+增加室内活动比例",
            "cost_reduction": "预计节省总预算15%",
            "user_impact": "minimal",
            "implementation": "调整1-2天行程+40%室内活动"
        })
        
        return counter_proposals

    def _select_optimal_solution(self, proposal_analysis, counter_proposals):
        """选择最优解决方案"""
        # 简单评分机制
        best_solution = None
        best_score = 0
        
        all_options = proposal_analysis + [{"proposal": cp} for cp in counter_proposals]
        
        for option in all_options:
            proposal = option.get("proposal", option)
            
            # 评分标准：成本效益 + 用户体验 + 可行性
            score = 0
            if option.get("recommendation") == "preferred":
                score += 3
            elif option.get("recommendation") == "acceptable":
                score += 2
            
            if option.get("feasibility") == "high":
                score += 2
            
            if option.get("user_impact") == "low":
                score += 2
            elif option.get("user_impact") == "medium":
                score += 1
            
            if score > best_score:
                best_score = score
                best_solution = proposal
        
        return {
            "selected_option": best_solution,
            "confidence": best_score / 7.0,
            "rationale": "基于成本效益、用户体验和可行性的综合评估"
        }

    def _get_budget_status(self, base_budget, budget_breakdown):
        """获取预算状态"""
        total_required = sum(budget_breakdown.values())
        min_viable_budget = 200  # 最低可行预算
        
        if base_budget < min_viable_budget:
            return "impossible"  # 预算太低，无法进行旅行
        elif total_required <= base_budget:
            return "sufficient"
        elif total_required <= base_budget * 1.1:
            return "tight"
        else:
            return "insufficient"

    def _analyze_budget_flexibility(self, budget_breakdown, user_requirements):
        """分析预算灵活性"""
        flexibility_score = 0.5  # 基础分数
        
        # 根据用户约束调整
        constraints = user_requirements.get("constraints", "")
        if "预算" in constraints or "budget" in constraints.lower():
            flexibility_score -= 0.2
        
        return {
            "flexibility_score": flexibility_score,
            "adjustable_categories": ["food", "activities"],
            "fixed_categories": ["accommodation", "transportation"],
            "max_adjustment_range": "±15%"
        }

    def _calculate_optimization_potential(self, budget_breakdown):
        """计算优化潜力"""
        # 简单计算：假设可以通过优化节省5-15%
        total_budget = sum(budget_breakdown.values())
        return {
            "potential_savings": total_budget * 0.1,
            "optimization_areas": ["住宿选择", "交通方式", "活动时段"],
            "confidence": 0.75
        }

    def _identify_cost_risks(self, budget_breakdown, user_requirements):
        """识别成本风险"""
        risks = []
        
        if budget_breakdown["emergency"] < sum(budget_breakdown.values()) * 0.05:
            risks.append({
                "type": "insufficient_emergency_fund",
                "severity": "medium",
                "description": "应急资金不足"
            })
        
        return risks

    def _suggest_collaboration_strategy(self, conflicts, budget_breakdown):
        """建议协作策略"""
        if conflicts:
            return {
                "strategy": "proactive_negotiation",
                "priority": ["cost_reduction", "user_satisfaction", "feasibility"],
                "collaboration_targets": ["WeatherAgent", "CoordinatorAgent"]
            }
        else:
            return {
                "strategy": "optimization_focus",
                "priority": ["cost_efficiency", "experience_quality"],
                "collaboration_targets": ["WeatherAgent"]
            }

    def _identify_overbudget_categories(self, budget_breakdown, base_budget):
        """识别超预算类别"""
        overbudget = []
        for category, amount in budget_breakdown.items():
            max_amount = base_budget * self.budget_categories[category]["max_ratio"]
            if amount > max_amount:
                overbudget.append(category)
        return overbudget

    def _handle_simple_budget_analysis(self, request_data):
        """处理简单预算分析"""
        budget_amount = 800  # 默认预算
        
        breakdown = self._analyze_budget_breakdown(budget_amount, request_data)
        
        return {
            "agent": "BudgetAgent",
            "budget_analysis": {
                "total_budget": budget_amount,
                "breakdown": breakdown,
                "recommendations": ["优先住宿安全", "预留应急资金", "灵活调整活动预算"]
            }
        }

    def initiate_direct_negotiation(self, target_agent: str, negotiation_topic: str, proposal: dict) -> str:
        """发起与其他Agent的直接协商"""
        if target_agent not in self.peer_agents:
            print(f"❌ 未知的目标Agent: {target_agent}")
            return ""

        negotiation_id = str(uuid.uuid4())[:8]

        # 构建A2A协商消息
        negotiation_message = {
            "type": "direct_negotiation_request",
            "negotiation_id": negotiation_id,
            "sender": "BudgetAgent",
            "topic": negotiation_topic,
            "proposal": proposal,
            "timestamp": datetime.now().isoformat()
        }

        try:
            # 直接发送A2A消息，不通过中介
            target_url = self.peer_agents[target_agent]
            payload = {
                "message": {
                    "content": {
                        "text": json.dumps(negotiation_message, ensure_ascii=False)
                    }
                }
            }

            print(f"🤝 BudgetAgent 直接向 {target_agent} 发起协商: {negotiation_topic}")
            response = requests.post(f"{target_url}/a2a/tasks/send", json=payload, timeout=15)

            if response.status_code == 200:
                # 记录协商会话
                self.active_direct_negotiations[negotiation_id] = {
                    "target": target_agent,
                    "topic": negotiation_topic,
                    "status": "active",
                    "proposal": proposal,
                    "start_time": datetime.now().isoformat()
                }
                return negotiation_id
            else:
                print(f"❌ 协商请求发送失败: {response.status_code}")
                return ""

        except Exception as e:
            print(f"❌ A2A协商通信失败: {e}")
            return ""

    def handle_direct_negotiation_request(self, sender: str, negotiation_data: dict) -> dict:
        """处理来自其他Agent的直接协商请求"""
        negotiation_id = negotiation_data.get("negotiation_id")
        topic = negotiation_data.get("topic")
        proposal = negotiation_data.get("proposal", {})

        print(f"📥 BudgetAgent 收到来自 {sender} 的直接协商: {topic}")

        # 评估协商提案
        evaluation = self._evaluate_peer_proposal(sender, topic, proposal)

        # 生成响应
        response = {
            "type": "direct_negotiation_response",
            "negotiation_id": negotiation_id,
            "sender": "BudgetAgent",
            "target": sender,
            "evaluation": evaluation
        }

        # 如果需要反提案
        if evaluation["status"] == "needs_modification":
            counter_proposal = self._generate_budget_counter_proposal(proposal, evaluation)
            response["counter_proposal"] = counter_proposal
            print(f"💡 BudgetAgent 向 {sender} 提出反提案")

        # 直接回复给发起方（真正的A2A）
        self._send_direct_response(sender, response)

        return response

    def _evaluate_peer_proposal(self, sender: str, topic: str, proposal: dict) -> dict:
        """评估来自对等Agent的提案"""
        if topic == "weather_cost_adjustment" and sender == "WeatherAgent":
            # WeatherAgent提议的成本调整
            extra_cost = proposal.get("extra_cost", 0)
            weather_condition = proposal.get("weather_condition", "")

            # 预算影响评估
            if extra_cost <= 100:  # 可接受范围
                return {
                    "status": "acceptable",
                    "confidence": 0.85,
                    "reasoning": f"额外成本{extra_cost}在可接受范围内",
                    "budget_adjustment": self._calculate_budget_reallocation(extra_cost)
                }
            elif extra_cost <= 200:  # 需要协商
                return {
                    "status": "needs_modification",
                    "confidence": 0.60,
                    "reasoning": f"成本较高，需要优化方案",
                    "suggested_reduction": extra_cost * 0.3
                }
            else:  # 不可接受
                return {
                    "status": "rejected",
                    "confidence": 0.20,
                    "reasoning": "成本过高，超出预算承受能力",
                    "max_acceptable": 200
                }

        # 默认评估
        return {
            "status": "needs_review",
            "confidence": 0.50,
            "reasoning": f"需要进一步分析{topic}提案"
        }

    def _send_direct_response(self, target_agent: str, response: dict):
        """直接发送响应给目标Agent"""
        if target_agent.lower().replace("agent", "") in self.peer_agents:
            agent_key = target_agent.lower().replace("agent", "")
            target_url = self.peer_agents[agent_key]

            payload = {
                "message": {
                    "content": {
                        "text": json.dumps(response, ensure_ascii=False)
                    }
                }
            }

            try:
                requests.post(f"{target_url}/a2a/tasks/send", json=payload, timeout=10)
                print(f"📤 BudgetAgent 直接回复 {target_agent}")
            except Exception as e:
                print(f"❌ 直接回复失败: {e}")

if __name__ == "__main__":
    agent_card = AgentCard(
        name="Budget Optimizer Agent",
        description="智能预算优化专家，提供预算分析、成本控制和Agent间预算协商服务",
        url="http://localhost:7007",
        version="1.0.0",
        skills=[
            AgentSkill(
                name="Budget Analysis",
                description="智能预算分析和分配优化"
            ),
            AgentSkill(
                name="Cost Negotiation", 
                description="与其他Agent协商成本控制方案"
            ),
            AgentSkill(
                name="Risk Assessment",
                description="识别和评估预算风险"
            ),
            AgentSkill(
                name="Alternative Planning",
                description="生成预算替代方案"
            )
        ]
    )

    budget_agent = BudgetAgent(agent_card)
    
    print("💰 BudgetAgent 启动于 http://localhost:7007")
    print("🤝 支持主动协商和冲突解决")
    run_server(budget_agent, host="0.0.0.0", port=7007)
